package org.example.sink;

import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.serialization.SimpleStringEncoder;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
import org.apache.flink.configuration.MemorySize;
import org.apache.flink.connector.datagen.source.DataGeneratorSource;
import org.apache.flink.connector.datagen.source.GeneratorFunction;
import org.apache.flink.core.fs.Path;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.filesystem.OutputFileConfig;
import org.apache.flink.streaming.api.functions.sink.filesystem.bucketassigners.DateTimeBucketAssigner;
import org.apache.flink.streaming.api.functions.sink.filesystem.rollingpolicies.DefaultRollingPolicy;

import java.time.Duration;

/**
 * @Author: tang
 * @Description:
 * @Date 2025/2/18 22:12
 */
public class FileSink {

    public static void main(String[] args) throws Exception{

        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();

        environment.setParallelism(2);
        environment.enableCheckpointing(1000, CheckpointingMode.EXACTLY_ONCE);

        DataGeneratorSource<Integer> generatorSource =
                new DataGeneratorSource<>((GeneratorFunction<Long, Integer>) Long::intValue,
                        Integer.MAX_VALUE, RateLimiterStrategy.perSecond(1), Types.INT);

        // 文件输出配置
        org.apache.flink.connector.file.sink.FileSink<String> fileSink = org.apache.flink.connector.file.sink.FileSink.<String>forRowFormat(
                        new Path("/Users/tanglonglong/Desktop"),
                        new SimpleStringEncoder<String>()
                )
                .withBucketAssigner(new DateTimeBucketAssigner<>("yyyy-MM-dd--HH"))
                .withOutputFileConfig(new OutputFileConfig("test", ".log"))
                // 滚动策略
                .withRollingPolicy(
                        DefaultRollingPolicy.builder()
                                .withRolloverInterval(
                                        Duration.ofSeconds(10)
                                )
                                .withMaxPartSize(new MemorySize(1024 * 1024))
                                .build()
                )
                .build();

        environment.fromSource(generatorSource, WatermarkStrategy.noWatermarks(), "gen-data")
                .map(String::valueOf)
                .sinkTo(fileSink);

        environment.execute();
    }

}
